import pandas as pd
import matplotlib.pyplot as plt

# 筛选出5个要描述的物种
species_to_change = ['Plankton', 'Small Fish', 'Medium Fish', 'Big Fish', 'Squid']

# 模拟物种的增长率
growth_rates = {
    'Plankton': 0.05,
    'Small Fish': 0.08,
    'Medium Fish': 0.1,
    'Big Fish': 0.03,
    'Squid': 0.06
}

# 模拟物种的死亡率
death_rates = {
    'Plankton': 0.04,
    'Small Fish': 0.09,
    'Medium Fish': 0.2,
    'Big Fish': 0.05,
    'Squid': 0.08
}

def species_change(marine_data,radiation_data):
    species_init={species: [] for species in species_to_change}
    species_changes={species: [] for species in species_to_change}
    # 创建时间序列
    # 转换时间列为日期时间格式
    timestamps = pd.to_datetime(radiation_data['timestamp'])
    
    # 使用线性回归模型对每个物种的数量进行描述
    for species in species_to_change:
        # 获取特定物种的数量数据
        species_quantity = marine_data[marine_data['species'] == species]['quantity'].values[0]
        # 获取特定时期的污染浓度数据
        radiation_concentration = radiation_data['concentration']
        # 模拟污染浓度对物种增长的影响
        effect_of_radiation = 1 - radiation_concentration*0.5  # 调整污染浓度对增长率的影响
        # 模拟污染浓度对物种死亡的影响
        death_effect_of_radiation = 1 + radiation_concentration*3

        # 计算每个时期物种的数量变化
        for i in range(len(radiation_concentration)):
            # 计算变化的物种数量
            # print(death_rates[species]*death_effect_of_radiation[i])
            # print(growth_rates[species] * effect_of_radiation[i])
            # 正常无辐射状态下的物种情况
            changed_quantity_normal = species_quantity * (1 + growth_rates[species]) * (1 - death_rates[species])
            # 有辐射后的情况
            changed_quantity = species_quantity * (1 + growth_rates[species] * effect_of_radiation[i]) * (1 - death_rates[species] * death_effect_of_radiation[i])
            
            species_init[species].append(changed_quantity_normal)
            species_changes[species].append(changed_quantity)
    
    # 绘制每个物种的数量变化图
    plt.figure(figsize=(10, 6))
    for species in species_to_change:
        # 不受辐射影响时生物数量的变化情况
        # plt.plot(timestamps, species_init[species], label=species+'-init')
        # 在辐射影响下生物数量的变化情况
        plt.plot(timestamps, species_changes[species], label=species+'-radiated')

        # 收集每个物种的数量变化数据到数据帧中
        species_data = pd.DataFrame({
            'Timestamp': timestamps,
            'Quantity_Changed': species_changes[species],
            'Radiated': [species + '-radiated'] * len(timestamps)  # 添加辐射标签
        })

        # 保存数据到 CSV 文件，文件名使用物种标签
        file_name = f'data_resource/species/{species}_data.csv'
        species_data.to_csv(file_name, index=False)

    plt.xlabel('Time')
    plt.ylabel('Quantity')
    plt.title('Changed Quantity of Species over Time')
    plt.legend()
    plt.xticks(rotation=45)
    plt.tight_layout()
    plt.show()

    return species_init,species_changes